from abc import ABC
import math

import tensorflow as tf


class SubCenterProduct(tf.keras.layers.Layer):
    """subCenter which can auto detect noisy sample"""
    """https://www.ecva.net/papers/eccv_2020/papers_ECCV/papers/123560715.pdf"""

    def __init__(self,
                 n_classes,
                 s=30.0,
                 m=0.5,
                 ls_eps=0.0,
                 k=3,
                 easy_margin=True):
        super(SubCenterProduct, self).__init__()
        self.s = s
        self.m = m
        # self.w = self.add_weight(name='W', shape=(in_dim, n_classes))
        self.easy_margin = easy_margin

        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        self.th = math.cos(math.pi - m)
        self.mm = math.sin(math.pi - m) * m
        self.n_classes = n_classes
        self.ls_eps = ls_eps  # 增加噪音，提高泛化
        self.k = k

    def build(self, input_shape):
        super(SubCenterProduct, self).build(input_shape[0])

        self.W = self.add_weight(
            name='W',
            shape=(int(input_shape[0][-1]), self.k, self.n_classes),
            initializer='glorot_uniform',
            dtype='float32',
            regularizer=tf.keras.regularizers.l2(0.0001),
            trainable=True)

    @tf.function
    def call(self, inputs, training=None, mask=None):
        raw, y = inputs
        y = tf.cast(y, dtype=tf.int32)
        cosine = tf.tensordot(
            tf.math.l2_normalize(raw, axis=1),
            tf.math.l2_normalize(self.W, axis=0),
            axes=1
        )
        cosine = tf.keras.layers.GlobalMaxPool1D()(cosine)
        sine = tf.math.sqrt(1.0 - tf.math.pow(cosine, 2))
        phi = cosine * self.cos_m - sine * self.sin_m
        """cos(m+delta)"""
        """sine = sqrt(1-cosine**2)"""
        if self.easy_margin:
            phi = tf.where(cosine > 0, phi, cosine)
        else:
            phi = tf.where(cosine > self.th, phi, cosine - self.mm)

        one_hot = tf.cast(
            tf.one_hot(y, depth=self.n_classes),
            dtype=cosine.dtype
        )

        if self.ls_eps > 0:
            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.n_classes

        output = one_hot * phi + (1 - one_hot) * cosine

        output = output * self.s

        return output


class ArcMarginProduct(tf.keras.layers.Layer):
    """cos(theta + m)"""

    def __init__(self,
                 n_classes,
                 s=30.0,
                 m=0.5,
                 ls_eps=0.0,
                 easy_margin=True):
        super(ArcMarginProduct, self).__init__()
        self.s = s
        self.m = m
        # self.w = self.add_weight(name='W', shape=(in_dim, n_classes))
        self.easy_margin = easy_margin

        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        self.th = math.cos(math.pi - m)
        self.mm = math.sin(math.pi - m) * m
        self.n_classes = n_classes
        self.ls_eps = ls_eps  # 增加噪音，提高泛化

    def build(self, input_shape):
        super(ArcMarginProduct, self).build(input_shape[0])

        self.W = self.add_weight(
            name='W',
            shape=(int(input_shape[0][-1]), self.n_classes),
            initializer='glorot_uniform',
            dtype='float32',
            regularizer=tf.keras.regularizers.l2(0.0001),
            trainable=True)

    @tf.function
    def call(self, inputs, training=None, mask=None):
        raw, y = inputs
        y = tf.cast(y, dtype=tf.int32)
        cosine = tf.matmul(
            tf.math.l2_normalize(raw, axis=1),
            tf.math.l2_normalize(self.W, axis=0)
        )
        sine = tf.math.sqrt(1.0 - tf.math.pow(cosine, 2))
        phi = cosine * self.cos_m - sine * self.sin_m
        """cos(m+delta)"""
        """sine = sqrt(1-cosine**2)"""
        if self.easy_margin:
            phi = tf.where(cosine > 0, phi, cosine)
        else:
            phi = tf.where(cosine > self.th, phi, cosine - self.mm)

        one_hot = tf.cast(
            tf.one_hot(y, depth=self.n_classes),
            dtype=cosine.dtype
        )

        if self.ls_eps > 0:
            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.n_classes

        output = one_hot * phi + (1 - one_hot) * cosine

        output = output * self.s

        return output


class CosMarginProduct(tf.keras.layers.Layer):
    """cos(delta)-m"""

    def __init__(self,
                 n_classes,
                 s=60,
                 m=0.5,
                 ls_eps=0.001):
        super(CosMarginProduct, self).__init__()
        self.m = m
        self.s = s
        self.n_classes = n_classes
        self.ls_eps = ls_eps

    def build(self, input_shape):
        super(CosMarginProduct, self).build(input_shape[0])

        self.W = self.add_weight(
            name='W',
            shape=(int(input_shape[0][-1]), self.n_classes),
            initializer='glorot_uniform',
            dtype='float32',
            regularizer=tf.keras.regularizers.l2(0.0001),
            trainable=True)

    @tf.function
    def call(self, inputs, training=None, mask=None):
        raw, y = inputs
        y = tf.cast(y, dtype=tf.int32)
        cosine = tf.matmul(
            tf.math.l2_normalize(raw, axis=1),
            tf.math.l2_normalize(self.W, axis=0)
        )

        phi = cosine - self.m  # cosine(x)-delta
        one_hot = tf.cast(tf.one_hot(y, depth=self.n_classes), dtype=cosine.dtype)

        if self.ls_eps > 0:
            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.n_classes

        output = one_hot * phi + (1 - one_hot) * cosine

        output = output * self.s

        return output


if __name__ == '__main__':
    pass
